import torch
import torch.nn as nn
from 词嵌入 import tokenizer, embedding
from 位置编码 import PositionEncoding

# 配置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义对话语言模型
class LanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        embed_dim = embedding.get_sentence_embedding_dimension()
        self.nhead = 8
        self.pe = PositionEncoding(embed_dim, device=device)
        # 获取次嵌入层并移动到设备上
        self.word_embeddings = embedding[0].auto_model.embeddings.word_embeddings
        self.word_embeddings = self.word_embeddings.to(device)
        # 初始化 transformer 模型
        self.transformer = nn.Transformer(
            d_model=embed_dim,
            nhead=self.nhead,
            num_encoder_layers=24,
            num_decoder_layers=24,
            dim_feedforward=embed_dim*4,
            batch_first=True,
            norm_first=True,        # True前层归一化【比后层归一化更稳定，当深度设置为10层以上的时候设置】
                                     # False后层归一化【深度较深的时候，可能导致梯度消失/爆炸】
            dropout=0.5
        )
        # 最后的线性层 分类器
        self.fc_out = nn.Linear(embed_dim, embedding.tokenizer.vocab_size).to(device)

    # ids 带批次的索引列表   [N, 50]
    def embedding(self, ids):
        return self.word_embeddings(ids)

    def forward(self, src, tgt, src_key_padding_mask, tgt_key_padding_mask):
        # 词嵌入 和 位置编码
        src = self.embedding(src)
        tgt = self.embedding(tgt)
        src = self.pe(src)
        tgt = self.pe(tgt)
        # 构造解码器的因果注意力掩码 【上三角矩阵】
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(
            tgt.size(1), device=device
        )
        tgt_mask = tgt_mask.expand(
            src.size(0)*self.nhead, -1, -1
        )
        # 调用transformer
        y = self.transformer(
            src,tgt,
            tgt_mask=tgt_mask,
            src_key_padding_mask=src_key_padding_mask,
            tgt_key_padding_mask=tgt_key_padding_mask,
            memory_key_padding_mask=src_key_padding_mask,
            tgt_is_causal=True      # 若为True，必须有tgt_mask 因果掩码
        )
        # 进行分类
        y = self.fc_out(y)
        return y

if __name__ == "__main__":
    from 处理数据集 import LanguageDataset
    from torch.utils.data import DataLoader
    # 初始化模型
    model = LanguageModel().to(device)
    # 初始化数据集 并 加载数据集
    ds = LanguageDataset()
    dl = DataLoader(ds, 5, shuffle=True)
    for i, (src, tgt, label) in enumerate(dl):
        src, src_key_padding_mask = tokenizer(src)
        tgt, tgt_key_padding_mask = tokenizer(tgt)
        label, _ = tokenizer(label)
        src, tgt, src_key_padding_mask, tgt_key_padding_mask = (src.to(device),
                                                                tgt.to(device),
                                                                src_key_padding_mask.to(device),
                                                                tgt_key_padding_mask.to(device))
        y = model(src, tgt, src_key_padding_mask, tgt_key_padding_mask)
        print(y.shape)

